use crate::benches::{
    MAX_LIST_ALLOWANCES, NUM_GET_BLOCKS, NUM_OPERATIONS, assert_has_num_balances,
    emulate_archive_blocks, icrc_transfer, mint_tokens, test_account, test_account_offset, upgrade,
};
use crate::{
    Access, LOG, icrc2_approve_not_async, icrc3_get_blocks, icrc103_get_allowances, init_state,
};
use assert_matches::assert_matches;
use canbench_rs::{BenchResult, bench};
use candid::{Nat, Principal};
use ic_icrc1_ledger::{FeatureFlags, InitArgs, InitArgsBuilder};
use ic_ledger_canister_core::archive::ArchiveOptions;
use icrc_ledger_types::icrc1::account::Account;
use icrc_ledger_types::icrc1::transfer::TransferArg;
use icrc_ledger_types::icrc2::approve::ApproveArgs;
use icrc_ledger_types::icrc3::blocks::GetBlocksRequest;
use icrc_ledger_types::icrc103::get_allowances::GetAllowancesArgs;

const MINTER_PRINCIPAL: Principal = Principal::from_slice(&[0_u8, 0, 0, 0, 2, 48, 0, 7, 1, 1]);

#[bench(raw)]
fn bench_upgrade_baseline() -> BenchResult {
    init_state(ckbtc_ledger_init_args_with_archive());
    assert_has_num_balances(0);

    canbench_rs::bench_fn(upgrade)
}

#[bench(raw)]
fn bench_icrc1_transfers() -> BenchResult {
    init_state(ckbtc_ledger_init_args_with_archive());
    let start_time = ic_cdk::api::time();
    let account_with_tokens = mint_tokens(MINTER_PRINCIPAL, u64::MAX);
    assert_has_num_balances(1);

    canbench_rs::bench_fn(|| {
        {
            let _p = canbench_rs::bench_scope("icrc1_transfer");
            for i in 0..NUM_OPERATIONS {
                let transfer = TransferArg {
                    from_subaccount: account_with_tokens.subaccount,
                    to: test_account(i),
                    created_at_time: Some(start_time + i as u64),
                    ..ckbtc_transfer()
                };
                let result = icrc_transfer(account_with_tokens.owner, None, transfer.clone());
                assert_matches!(result, Ok(_));
                emulate_archive_blocks::<Access>(&LOG);
            }
            assert_has_num_balances(NUM_OPERATIONS + 2);
        }
        {
            let _p = canbench_rs::bench_scope("icrc2_approve");
            for i in 0..NUM_OPERATIONS {
                let approve = ApproveArgs {
                    from_subaccount: account_with_tokens.subaccount,
                    spender: test_account(i),
                    created_at_time: Some(start_time + i as u64),
                    amount: u64::MAX.into(),
                    expected_allowance: Some(0u64.into()),
                    expires_at: Some(u64::MAX),
                    fee: None,
                    memo: Some(MEMO.to_vec().into()),
                };
                let result = icrc2_approve_not_async(account_with_tokens.owner, approve.clone());
                assert_matches!(result, Ok(_));
                emulate_archive_blocks::<Access>(&LOG);
            }
        }
        {
            let _p = canbench_rs::bench_scope("icrc103_get_allowances");
            let list_args = GetAllowancesArgs {
                from_account: Some(account_with_tokens),
                prev_spender: None,
                take: None,
            };
            let allowances = icrc103_get_allowances(list_args).expect("failed to list allowances");
            assert_eq!(allowances.len(), MAX_LIST_ALLOWANCES);
        }
        {
            let _p = canbench_rs::bench_scope("icrc2_transfer_from");
            for i in 0..NUM_OPERATIONS {
                let spender = test_account(i);
                let transfer = TransferArg {
                    from_subaccount: account_with_tokens.subaccount,
                    to: test_account_offset(i),
                    created_at_time: Some(start_time + i as u64),
                    ..ckbtc_transfer()
                };
                let result =
                    icrc_transfer(account_with_tokens.owner, Some(spender), transfer.clone());
                assert_matches!(result, Ok(_));
                emulate_archive_blocks::<Access>(&LOG);
            }
            assert_has_num_balances(2 * NUM_OPERATIONS + 2);
        }
        for i in 0..NUM_GET_BLOCKS {
            let spender = test_account(i);
            let transfer = TransferArg {
                from_subaccount: account_with_tokens.subaccount,
                to: test_account_offset(i),
                created_at_time: Some(1_000_000_000 + start_time + i as u64),
                ..ckbtc_transfer()
            };
            let result = icrc_transfer(account_with_tokens.owner, Some(spender), transfer.clone());
            assert_matches!(result, Ok(_));
        }
        {
            let req = GetBlocksRequest {
                start: Nat::from(3 * NUM_OPERATIONS),
                length: Nat::from(NUM_GET_BLOCKS),
            };
            let _p = canbench_rs::bench_scope("icrc3_get_blocks");
            let blocks_res = icrc3_get_blocks(vec![req]);
            assert_eq!(blocks_res.blocks.len(), NUM_GET_BLOCKS as usize);
        }
        upgrade();
    })
}

fn ckbtc_ledger_init_args_with_archive() -> InitArgs {
    let minter_principal = Principal::from_text("mqygn-kiaaa-aaaar-qaadq-cai").unwrap();
    assert_eq!(minter_principal, MINTER_PRINCIPAL);
    let nns_root_principal = Principal::from_text("r7inp-6aaaa-aaaaa-aaabq-cai").unwrap();
    InitArgsBuilder::for_tests()
        .with_minting_account(minter_principal)
        .with_fee_collector_account(Account {
            owner: minter_principal,
            subaccount: Some([
                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                0, 0, 0x0f, 0xee,
            ]),
        })
        .with_decimals(8)
        .with_max_memo_length(80)
        .with_transfer_fee(10_u64)
        .with_token_symbol("ckBTC")
        .with_token_name("ckBTC")
        .with_feature_flags(FeatureFlags { icrc2: true })
        .with_metadata_entry("icrc1:logo", "")
        .with_archive_options(ArchiveOptions {
            trigger_threshold: 2_000,
            num_blocks_to_archive: 1_0000,
            node_max_memory_size_bytes: Some(3_221_225_472),
            max_message_size_bytes: None,
            controller_id: nns_root_principal.into(),
            more_controller_ids: None,
            cycles_for_archive_creation: Some(100_000_000_000_000),
            max_transactions_per_response: None,
        })
        .build()
}

const MEMO: [u8; 41] = [
    0x82_u8, 0x00, 0x83, 0x58, 0x20, 0x18, 0x19, 0xcc, 0xd2, 0x28, 0xad, 0x2e, 0x83, 0xc6, 0xc8,
    0x63, 0x99, 0xa0, 0xd7, 0xd0, 0x2e, 0xe9, 0x75, 0x96, 0x95, 0x86, 0xf3, 0x47, 0x85, 0xf6, 0xaf,
    0x99, 0x00, 0x1e, 0x08, 0x8b, 0xa0, 0x02, 0x19, 0x07, 0xd0,
];

/// ckBTC ledger transaction 1604556
fn ckbtc_transfer() -> TransferArg {
    TransferArg {
        from_subaccount: None,
        to: Account {
            owner: Principal::from_text(
                "rgwvx-m5lab-jv5b7-n6tc6-agk3j-fle2q-ndsbh-gh2wg-dxphx-oyk3i-vqe",
            )
            .unwrap(),
            subaccount: None,
        },
        fee: None,
        created_at_time: None,
        memo: Some(MEMO.to_vec().into()),
        amount: 167_708_u32.into(),
    }
}
